"""
model_config.py
---------------
定义与模型结构相关的配置，包括：
- 模型类型（内置 / 自定义）
- 模型架构名称或自定义模型脚本路径
- 模型参数（类别数、输入维度等）
- 是否使用预训练权重、权重路径
"""

from dataclasses import dataclass
from typing import Optional, Dict

@dataclass
class ModelConfig:
    # === 基础信息 ===
    model_type: str = "builtin"         # 模型类型："builtin" 或 "custom"
    model_name: str = "resnet18"        # 内置模型名称（当 model_type="builtin" 时使用）
    num_classes: int = 10               # 输出类别数
    input_channels: int = 3             # 输入通道数（RGB=3, 灰度=1）
    dropout: float = 0.5                # Dropout 比例，用于防止过拟合（0 表示不使用 Dropout，值越大丢弃的神经元越多）

    # === 预训练权重 ===
    pretrained: bool = False            # 是否使用预训练模型
    pretrained_path: Optional[str] = None  # 自定义预训练权重路径 (.pth 或 .pt)

    # === 自定义模型 ===
    custom_model_path: Optional[str] = None
    """
    指定本地 Python 模型定义文件路径，如:
    './models/my_custom_model.py'
    当 model_type='custom' 时生效。
    """

    custom_model_class: Optional[str] = None
    """
    指定自定义模型类名，例如 'MyNet'。
    系统会通过 importlib 动态导入该类。
    """

    # === 额外模型参数 ===
    model_args: Optional[Dict] = None
    """
    模型构造参数（用于自定义模型初始化）。
    例如: {'hidden_dim': 256, 'dropout': 0.5}
    """

    # === 设备 & 其他控制 ===
    device: str = "cuda"                # 模型运行设备
    use_half_precision: bool = False    # 是否启用混合精度训练

    def summary(self):
        """打印模型配置信息"""
        print("🧠 模型配置")
        print(f"  模型类型: {self.model_type}")
        print(f"  模型名称: {self.model_name}")
        print(f"  类别数: {self.num_classes}")
        print(f"  输入通道: {self.input_channels}")
        if self.model_type == "custom":
            print(f"  自定义模型路径: {self.custom_model_path}")
            print(f"  自定义模型类名: {self.custom_model_class}")
        if self.pretrained:
            print(f"  使用预训练模型: {self.pretrained_path or '默认预训练权重'}")
        print(f"  设备: {self.device}, 混合精度: {self.use_half_precision}")
